import  torch
from torch import nn, optim
from torch.nn import functional as F
import torchvision
from    torch.utils.data import DataLoader
from    torchvision import datasets
from    torchvision import transforms

class BasicBlock(nn.Module):
    def __init__(self, input_filter_num, output_filter_num, stride=1):
        super(BasicBlock, self).__init__()
        # 第一个卷积单元
        self.conv1 = nn.Conv2d(input_filter_num, output_filter_num, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(output_filter_num)
        self.relu = nn.ReLU()
        # 第二个卷积单元
        self.conv2 = nn.Conv2d(output_filter_num, output_filter_num, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(output_filter_num)
        # shape匹配，直接短接
        if input_filter_num != output_filter_num:
            self.downsample = nn.Sequential(nn.Conv2d(input_filter_num, output_filter_num, kernel_size=1, stride=stride),
                                            nn.BatchNorm2d(output_filter_num))
        else:
            self.downsample = lambda x : x

    def forward(self, inputs):
        x = self.relu(self.bn1(self.conv1(inputs)))
        out1 = self.bn2(self.conv2(x))

        identity = self.downsample(inputs)
        output = out1 + identity

        return output

class ResNet(nn.Module):
    def __init__(self, layer_dims, num_class=100): # layer_Dim[2, 2, 2, 2]
        super(ResNet, self).__init__()

        self.stem = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
                                  nn.BatchNorm2d(64),
                                  nn.ReLU(),
                                  nn.MaxPool2d(kernel_size=3, stride=1, padding=1))
        self.layer1 = self.build_resblock(64, 64, layer_dims[0])
        self.layer2 = self.build_resblock(64, 128, layer_dims[1], stride=2)
        self.layer3 = self.build_resblock(128, 256, layer_dims[2], stride=2)
        self.layer4 = self.build_resblock(256, 512, layer_dims[3], stride=2)

        self.fc = nn.Linear(512, num_class)

    def forward(self, inputs):
        x = self.stem(inputs)
        #print('stem:', x.shape)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        #print('after 4conv:', x.shape)
        # [b,c, 1, 1]
        x = F.adaptive_avg_pool2d(x, (1,1))
        #print('gobal_avg_pool:', x.shape)
        # [b,c]
        x = x.view(x.shape[0], -1)
        #print('reshape:', x.shape)
        #[b, 100]
        x = self.fc(x)
        #print('out:', x.shape)

        return x


    def build_resblock(self, input_filter_num, output_filter_num, blocks, stride=1):
        '''
        构建resblock
        :param input_filter_num: 输入channel数
        :param output_filter_num: 输出的channel数
        :param blocks: 一个resblock包涵几个basicblock
        :param stride: 卷积步长，实现下采样
        :return: 返回resblock
        '''
        res_block = nn.Sequential()
        # 只有第一个basicblock下采样
        res_block.add_module('model', BasicBlock(input_filter_num, output_filter_num, stride))

        for i in range(1, blocks):
            res_block.add_module('model%d'%i, BasicBlock(output_filter_num, output_filter_num, 1))

        return res_block


def main():

    data = torch.rand(128, 3, 32, 32)
    model = ResNet([3, 4, 6, 3])
    model(data)

def ResNet18():
    return ResNet([2, 2, 2, 2])

def ResNet34():
    return ResNet([3, 4, 6, 3])

if __name__ == '__main__':
    main()